import fmeobjects
import math
import time
import numpy as np
import pickle


def calculate_distance(point1, point2):
    # 转换为 NumPy 数组
    point1 = np.array(point1)
    point2 = np.array(point2)

    # 计算坐标差的平方
    diff = point2 - point1
    squared_diff = np.square(diff)

    # 求和并取平方根
    distance = np.sqrt(np.sum(squared_diff))
 
    return distance

# 计算向量上，距离终点一定距离的点坐标
# 如果向量为A->B, 点在AB上，距离B一定距离
# direction_vector 方向向量
# point 点坐标
# distance 距离
def calculate_point_on_vector(direction_vector, point, distance):
    length = np.linalg.norm(direction_vector)  # 方向向量的长度
    scaled_vector = (distance / length) * direction_vector  # 缩放后的方向向量
    return np.array([point[0] - scaled_vector[0], point[1] - scaled_vector[1], point[2] - scaled_vector[2]])  # 计算目标点坐标

# 计算向量A以向量B为轴旋转一定角度后的新向量
# vector 为要旋转的向量
# axis 旋转轴向量
# angle 旋转角度（弧度）
def rotate_vector_around_axis(vector, axis, angle):
    # 归一化向量
    vector = vector / np.linalg.norm(vector)
    axis = axis / np.linalg.norm(axis)

    # 计算旋转四元数
    q = quaternion_from_axis_angle(axis, angle)

    # 执行四元数旋转
    rotated_vector = quaternion_rotate_vector(q, vector)

    return rotated_vector


def quaternion_from_axis_angle(axis, angle):
    # 计算旋转四元数
    angle /= 2
    axis = axis / np.linalg.norm(axis)
    q = np.array([np.cos(angle), axis[0] * np.sin(angle), axis[1] * np.sin(angle), axis[2] * np.sin(angle)])

    return q


def quaternion_rotate_vector(q, v):
    # 执行四元数旋转
    q_conj = np.array([q[0], -q[1], -q[2], -q[3]])
    rotated_vector = quaternion_multiply(quaternion_multiply(q, np.concatenate(([0], v))), q_conj)[1:]

    return rotated_vector


def quaternion_multiply(q1, q2):
    # 四元数乘法
    w1, x1, y1, z1 = q1
    w2, x2, y2, z2 = q2
    w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
    x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
    y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
    z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2

    return np.array([w, x, y, z])


def calculate_vertex(normal, center, pipe_shape, pipe_width, pipe_height, num_interpolation_points):
    distance = 0
    angle_array = []
    if pipe_shape == 1:  # 圆管
        # num_interpolation_points = 12
        distance = pipe_height
        # angle_array = [0, np.pi / 6, np.pi / 3, np.pi / 2, 2 * np.pi / 3, 5 * np.pi / 6, np.pi, 7 * np.pi / 6,
        #                4 * np.pi / 3, 3 * np.pi / 2, 5 * np.pi / 3, 11 * np.pi / 6]
        for i in range(num_interpolation_points):
            angle_array.append(2 * i * np.pi / num_interpolation_points)
    else:  # 方管
        distance = math.sqrt(math.pow(pipe_width / 2, 2) + math.pow(pipe_height / 2, 2))
        angle_tmp = math.atan(pipe_width / pipe_height)
        angle_array = [angle_tmp, np.pi - angle_tmp, np.pi + angle_tmp, 2 * np.pi - angle_tmp]

    # 旋转轴为 normal
    # 首先计算要旋转的向量
    v1 = np.cross(normal, np.array([0, 0, 1]))
    vector_rotate = np.cross(normal, v1)

    result_array = []
    for angle in angle_array:
        vector_tmp = rotate_vector_around_axis(vector_rotate, normal, angle)
        u_tmp = vector_tmp / np.linalg.norm(vector_tmp)
        pt = center + u_tmp * distance
        result_array.append((round(pt[0], 3), round(pt[1], 3), round(pt[2], 3)))

    return result_array


def divide_arc(center, start_point, end_point, n):
    v1 = start_point - center
    v2 = end_point - center

    r = np.linalg.norm(v1)

    u1 = v1 / np.linalg.norm(v1)
    u2 = v2 / np.linalg.norm(v2)

    # 计算旋转角度
    theta = np.arccos(np.dot(u1, u2))

    rotation_axis = np.cross(u1, u2)
    rotation_axis /= np.linalg.norm(rotation_axis)

    #rotation_matrix = rotation_matrix_from_axis_angle(rotation_axis, theta)

    points = []
    for i in range(n):
        angle = i * theta / (n - 1)
        rotated_point = rotate_around_axis(start_point, center, rotation_axis, angle)
        points.append(rotated_point.tolist())

    return points


# 求点绕向量旋转一定角度后的新的点位坐标
def rotate_around_axis(point, axis_origin, axis_direction, angle):
    radius = calculate_distance(point,axis_origin)
    start_vector = np.array(point-axis_origin)
    end_vector = rotate_vector_around_axis(start_vector,axis_direction,angle)
    rotated_point = calculate_point_on_vector(-end_vector,axis_origin,radius)
    return rotated_point


def find_intersection(line1_point, line1_direction, line2_point, line2_direction):
    A = np.vstack((line1_direction, -line2_direction)).T
    b = line2_point - line1_point

    try:
        t, s = np.linalg.lstsq(A, b, rcond=None)[0]
    except np.linalg.LinAlgError:
        return None

    intersection_point = line1_point + t * line1_direction
    # intersection_point = line2_point + s * line2_direction
    return intersection_point


# 在直线1，直线2的交点处计算相关参数
def calculate_pipe_params(start_point_line1, end_point_line1, start_point_line2, end_point_line2, turning_radius,
                          num_segments_on_arc, num_segments_per_circle_section, pipe_shape, pipe_width_line1,
                          pipe_height_line1, pipe_width_line2,
                          pipe_height_line2):
    vector_line1 = end_point_line1 - start_point_line1
    vector_line2 = start_point_line2 - end_point_line2
    tmp_end_point_line1 = calculate_point_on_vector(vector_line1, end_point_line1, turning_radius)
    tmp_start_point_line2 = calculate_point_on_vector(vector_line2, start_point_line2, turning_radius)
    plane_normal = np.cross(vector_line1, vector_line2)
    vector_line1_vertical = np.cross(plane_normal, vector_line1)
    vector_line2_vertical = np.cross(plane_normal, vector_line2)
    circle_center = find_intersection(tmp_end_point_line1, vector_line1_vertical, tmp_start_point_line2,
                                      vector_line2_vertical)

    arc_segments = divide_arc(circle_center, tmp_end_point_line1, tmp_start_point_line2, num_segments_on_arc)
    elbow_points_vertices = []

    for i in range(num_segments_on_arc):
        segment = arc_segments[i]
        tmp_pipe_width = 0
        tmp_pipe_height = 0
        if pipe_width_line1 == pipe_width_line2 and pipe_height_line1 == pipe_height_line2:
            tmp_pipe_width = pipe_width_line1
            tmp_pipe_height = pipe_height_line1
        else:
            tmp_pipe_width = pipe_width_line1 + i * (pipe_width_line2 - pipe_width_line1) / num_segments_on_arc
            tmp_pipe_height = pipe_height_line1 + i * (pipe_height_line2 - pipe_height_line1) / num_segments_on_arc
        tmp_vector_segment = np.cross(plane_normal, circle_center - segment)
        elbow_point_vertices = calculate_vertex(tmp_vector_segment, segment, pipe_shape, tmp_pipe_width,
                                                tmp_pipe_height, num_segments_per_circle_section)
        elbow_points_vertices.append(elbow_point_vertices)
    # return arc_segments, start_point_line1, tmp_end_point_line1, tmp_start_point_line2, end_point_line2
    return elbow_points_vertices


# class FeatureProcessor(object):
#     """Template Class Interface:
#     When using this class, make sure its name is set as the value of the 'Class
#     to Process Features' transformer parameter.
#     """

#     def __init__(self):
#         """Base constructor for class members."""
#         pass

#     def input(self, feature):
#         """This method is called for each FME Feature entering the
#         PythonCaller. If knowledge of all input Features is not required for
#         processing, then the processed Feature can be emitted from this method
#         through self.pyoutput(). Otherwise, the input FME Feature should be
#         cached to a list class member and processed in process_group() when
#         'Group by' attributes(s) are specified, or the close() method.

#         :param fmeobjects.FMEFeature feature: FME Feature entering the
#             transformer.
#         """
#         turning_radius = 0.3
#         num_segments_on_arc = 10
#         pipe_shape = feature.getAttribute("line{0}.PipeShape")
#         num_segments_per_section = (12 if pipe_shape == 1 else 4)
#         feature.setAttribute("num_segments_per_section",num_segments_per_section)
#         A, B, C, D = None, None, None, None

#         if feature.getAttribute("pcode") == feature.getAttribute("line{0}.fircode"):
#             A = np.array([feature.getAttribute("line{0}.secx"), feature.getAttribute("line{0}.secy"),
#                           feature.getAttribute("line{0}.secz")])
#             B = np.array([feature.getAttribute("line{0}.firx"), feature.getAttribute("line{0}.firy"),
#                           feature.getAttribute("line{0}.firz")])
#         elif feature.getAttribute("pcode") == feature.getAttribute("line{0}.seccode"):
#             A = np.array([feature.getAttribute("line{0}.firx"), feature.getAttribute("line{0}.firy"),
#                           feature.getAttribute("line{0}.firz")])
#             B = np.array([feature.getAttribute("line{0}.secx"), feature.getAttribute("line{0}.secy"),
#                           feature.getAttribute("line{0}.secz")])

#         if feature.getAttribute("pcode") == feature.getAttribute("line{1}.fircode"):
#             C = np.array([feature.getAttribute("line{1}.firx"), feature.getAttribute("line{1}.firy"),
#                           feature.getAttribute("line{1}.firz")])
#             D = np.array([feature.getAttribute("line{1}.secx"), feature.getAttribute("line{1}.secy"),
#                           feature.getAttribute("line{1}.secz")])

#         elif feature.getAttribute("pcode") == feature.getAttribute("line{1}.seccode"):
#             C = np.array([feature.getAttribute("line{1}.secx"), feature.getAttribute("line{1}.secy"),
#                           feature.getAttribute("line{1}.secz")])
#             D = np.array([feature.getAttribute("line{1}.firx"), feature.getAttribute("line{1}.firy"),
#                           feature.getAttribute("line{1}.firz")])
#         #print(f'A->D: {A} {B} {C} {D}')

#         arc_segments = calculate_pipe_params(A, B, C, D, turning_radius, num_segments_on_arc,
#                                              num_segments_per_section, pipe_shape,
#                                              feature.getAttribute("line{0}.Width")/1000.0,
#                                              feature.getAttribute("line{0}.Height")/1000.0,
#                                              feature.getAttribute("line{1}.Width")/1000.0,
#                                              feature.getAttribute("line{1}.Height")/1000.0)
#         vertices = np.concatenate(arc_segments)
#         newFeature = fmeobjects.FMEFeature()
#         mesh = fmeobjects.FMEMesh()
#         for vertex in vertices:
#             mesh.appendVertex(vertex[0], vertex[1], vertex[2])

#         for i in range(len(arc_segments)):
#             if i == 0:  # 构建起始截面
#                 tmp_index_array = []
#                 for k in range(num_segments_per_section):
#                     tmp_index_array.append(num_segments_per_section - k)
#                 mesh.addMeshPart(0, tmp_index_array, None, None)
#             elif i > 0:
#                 for j in range(num_segments_per_section):
#                     if j < num_segments_per_section - 1:
#                         mesh.addMeshPart(0, [(i - 1) * num_segments_per_section + j,
#                                              (i - 1) * num_segments_per_section + j + 1,
#                                              i * num_segments_per_section + j + 1,
#                                              i * num_segments_per_section + j], None, None)
#                     else:
#                         mesh.addMeshPart(0, [(i - 1) * num_segments_per_section + j,
#                                              (i - 1) * num_segments_per_section,
#                                              i * num_segments_per_section,
#                                              i * num_segments_per_section + j], None, None)
#                 if i == len(arc_segments[0]) - 1:
#                     tmp_index_array = []
#                     for k in range(num_segments_per_section):
#                         tmp_index_array.append(i * num_segments_per_section + k)
#                     mesh.addMeshPart(0, tmp_index_array, None, None)
#         feature.setGeometry(mesh)
#         self.pyoutput(feature)

# def close(self):
#     """This method is called once all the FME Features have been processed
#     from input().
#     """
#     pass


# def process_group(self):
#     """When 'Group By' attribute(s) are specified, this method is called
#     once all the FME Features in a current group have been sent to input().

#     FME Features sent to input() should generally be cached for group-by
#     processing in this method when knowledge of all Features is required.
#     The resulting Feature(s) from the group-by processing should be emitted
#     through self.pyoutput().

#     FME will continue calling input() a number of times followed
#     by process_group() for each 'Group By' attribute, so this
#     implementation should reset any class members for the next group.
#     """
#     pass


# def has_support_for(self, support_type):
#     """This method returns whether this PythonCaller supports a certain type.
#     The only supported type is fmeobjects.FME_SUPPORT_FEATURE_TABLE_SHIM.

#     :param int support_type: The support type being queried.
#     :returns: True if the passed in support type is supported.
#     :rtype: bool
#     """
#     if support_type == fmeobjects.FME_SUPPORT_FEATURE_TABLE_SHIM:
#         # If this is set to return True, FME will pass features to the input() method that
#         # come from a feature table object. This allows for significant performance gains
#         # when processing large numbers of features.
#         # To enable this, the following conditions must be met:
#         #   1) features passed into the input() method cannot be copied or cached for later use
#         #   2) features cannot be read or modified after being passed to self.pyoutput()
#         #   3) Group Processing must not be enabled
#         # Violations will cause undefined behavior.
#         return False

#     return False
